# 导入必要模块
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# 第1部分：读取数据
file_path = "Chip_Test_Data.csv"
data = pd.read_csv(file_path)

# 给 x 和 y 赋值
y = data.loc[:, 'pass_or_not']
x1 = data.loc[:, 'test1']
x2 = data.loc[:, 'test2']

# 手动创建高阶特征字典
x = {
    'x1': x1,
    'x2': x2,
    'x1^2': x1 ** 2,
    'x2^2': x2 ** 2,
    'x1*x2': x1 * x2,
    'x1^3': x1 ** 3,
    'x2^3': x2 ** 3,
    'x1^2*x2': (x1 ** 2) * x2,
    'x1*x2^2': x1 * (x2 ** 2),
}

# 将特征字典转换为 DataFrame
x = pd.DataFrame(x)

# 第3部分：训练逻辑回归模型
model = LogisticRegression(max_iter=10000)
model.fit(x, y)

# 获取模型的系数与截距
theta = model.coef_[0]
theta0 = model.intercept_[0]

# 获取预测值并计算准确率
prediction = model.predict(x)
accuracy = accuracy_score(y, prediction)
print(f"Accuracy: {accuracy}")

# 第4部分：生成网格，计算决策边界
# 定义 x1 和 x2 的范围，用于绘制网格决策边界
x1_min, x1_max = x1.min() - 1, x1.max() + 1
x2_min, x2_max = x2.min() - 1, x2.max() + 1

# 生成网格数据
xx1, xx2 = np.meshgrid(np.linspace(x1_min, x1_max, 500),
                       np.linspace(x2_min, x2_max, 500))

# 使用生成网格计算多项式特征
grid_x = {
    'x1': xx1.ravel(),
    'x2': xx2.ravel(),
    'x1^2': xx1.ravel() ** 2,
    'x2^2': xx2.ravel() ** 2,
    'x1*x2': xx1.ravel() * xx2.ravel(),
    'x1^3': xx1.ravel() ** 3,
    'x2^3': xx2.ravel() ** 3,
    'x1^2*x2': (xx1.ravel() ** 2) * xx2.ravel(),
    'x1*x2^2': xx1.ravel() * (xx2.ravel() ** 2),
}

grid_x = pd.DataFrame(grid_x)

# 计算每个网格点的预测值
z = model.predict(grid_x)
z = z.reshape(xx1.shape)

# 第5部分：可视化
# 绘制样本点
class0 = (y == 0)
class1 = (y == 1)

plt.figure(figsize=(8, 6))

# 绘制 Not Pass 和 Pass 样本点
plt.scatter(x1[class0], x2[class0], c='r', label='Not Pass', marker='o')
plt.scatter(x1[class1], x2[class1], c='b', label='Pass', marker='x')

# 绘制决策边界
plt.contour(xx1, xx2, z, levels=[0.5], colors='g')

# 添加标签和标题
plt.xlabel('Test 1 Score')
plt.ylabel('Test 2 Score')
plt.legend()
plt.title('Polynomial Logistic Regression (Degree=3)')
plt.show()
